#include "StdAfx.h"
#include "Process.h"
#include <Tlhelp32.h>
#include "DllInjectorMgmt.h"
#include <Psapi.h>
#pragma comment(lib,"Psapi.lib")

std::vector<ProcInfo > CProcess::m_ProcList(0);
HANDLE CProcess::m_MonitorThreadHandle=0;
BOOL CProcess::m_isMoning=FALSE;
CRITICAL_SECTION CProcess::m_cs={0};
IProcessEvents *CProcess::m_IEvents=NULL;



BOOL CProcess::IsProcessInList(LPSTR strFilePath)
{
	Enter();
	BOOL retval=FALSE;
	int iCount=m_ProcList.size();
	for (int i=0;i<iCount;i++)
	{
		if (m_ProcList[i].FilePath[0]!=0
			&&strFilePath[0]!=0)
		{
			if (strcmp(strFilePath,m_ProcList[i].FilePath))
			{
				retval=TRUE;
				break;
			}
		}
	}
	Leave();
	return retval;
}

BOOL CProcess::IsProcessInList( DWORD pid ,BOOL isNeedActive)
{
	Enter();
	BOOL retval=FALSE;
	int iCount=m_ProcList.size();
	for (int i=0;i<iCount;i++)
	{
		if (m_ProcList[i].pid==pid && pid!=0)
		{
			if (isNeedActive==TRUE)
			{
				m_ProcList[i].isActive=TRUE;
			}
			retval=TRUE;
			break;
		}
	}
	Leave();
	return retval;
}

int CProcess::GetProcListCount()
{
	return m_ProcList.size();
}

BOOL CProcess::GetProcInfo(int nID,ProcInfo &pi)
{
	Enter();
	BOOL retval=FALSE;
	int iCount=m_ProcList.size();
	if (nID<iCount)
	{
		pi=m_ProcList[nID];
		retval=TRUE;
	}
	Leave();
	return retval;
}

void CProcess::RefreshProcessList()
{
	Enter();
	PROCESSENTRY32 pe32={0};
	pe32.dwSize = sizeof(PROCESSENTRY32);
	HANDLE hTool32 = CreateToolhelp32Snapshot(TH32CS_SNAPALL, NULL); //Create the snapshot
	BOOL bProcess = Process32First(hTool32, &pe32); //Call this and see if its valid
	if(bProcess == TRUE)
	{
		UnActiveProcess();
		while((Process32Next(hTool32, &pe32)) == TRUE)
		{
			if (IsProcessInList(pe32.th32ProcessID)==FALSE)
			{
				ProcInfo pi={0};
				HANDLE hProcess=OpenProcess(PROCESS_QUERY_INFORMATION|PROCESS_VM_READ,FALSE,pe32.th32ProcessID);
				if (hProcess!=NULL)
				{
					HMODULE hPMod=0;
					DWORD cbNeeded=0;
					if (EnumProcessModulesEx(hProcess,&hPMod,sizeof(HMODULE),&cbNeeded,LIST_MODULES_ALL))
					{
						GetModuleFileNameEx(hProcess,hPMod,pi.FilePath,MAX_PATH);
					}
					CloseHandle(hProcess);
				}
				if (strlen(pi.FilePath)==0)
				{
					sprintf_s(pi.FilePath,MAX_PATH,"%s",pe32.szExeFile);
				}

				InjectDll(pe32.th32ProcessID,"detoured.dll");
				InjectDll(pe32.th32ProcessID,"APIHook.dll");
				
				pi.pid=pe32.th32ProcessID;
				pi.isActive=TRUE;
				m_ProcList.push_back(pi);

				 
				PostBack(event_AddNew,&pi,sizeof(ProcInfo));
			}
		}
	}


	int iCount=m_ProcList.size();
	for (int i=0;i<iCount;i++)
	{
		if (m_ProcList[i].isActive==FALSE)
		{
			PostBack(event_DelProc,&m_ProcList[i],sizeof(ProcInfo));
			m_ProcList.erase(m_ProcList.begin()+i);
			break;
		}
	}

	Leave();
}

void CProcess::StartMonitor()
{
	Enter();
	m_MonitorThreadHandle=CreateThread(NULL,NULL,MonitorThread,NULL,NULL,NULL);
	if (m_MonitorThreadHandle!=NULL)
	{
		m_isMoning=TRUE;
	}
	Leave();
}

void CProcess::StopMonitor()
{
	Enter();
	if (m_MonitorThreadHandle!=NULL)
	{
		m_isMoning=TRUE;
		Sleep(1);
		CloseHandle(m_MonitorThreadHandle);
		m_MonitorThreadHandle=NULL;
		m_ProcList.clear();
	}
	Leave();
}

DWORD WINAPI CProcess::MonitorThread( LPVOID lpParameter )
{
	while(m_isMoning)
	{
		RefreshProcessList();
		Sleep(500);
	}
	return S_OK;
}

inline void CProcess::Enter()
{
	EnterCriticalSection(&m_cs);
}

inline void CProcess::Leave()
{
	LeaveCriticalSection(&m_cs);
}

void CProcess::Initialize()
{
	InitializeCriticalSection(&m_cs);
	m_ProcList.clear();
}

void CProcess::Release()
{
	m_ProcList.clear();
	DeleteCriticalSection(&m_cs);
	
}

inline void CProcess::PostBack( ProcessEvents eEvents ,LPVOID postBuf,DWORD bufSize)
{
	if (m_IEvents!=NULL)
	{
		m_IEvents->OnProcessEvents(eEvents,postBuf,bufSize);
	}
}

void CProcess::SetPostBack( IProcessEvents *IEvents )
{
	m_IEvents=IEvents;
}

void CProcess::UnActiveProcess()
{
	Enter();
	int iCount=m_ProcList.size();
	for (int i=0;i<iCount;i++)
	{
		m_ProcList[i].isActive=FALSE;
	}
	Leave();
}
void CProcess::InjectDll(DWORD PID,LPSTR DllPath)
{
	CDllInjector::Reset();
	CDllInjector::OpenProcess(PID,TRUE);
	CDllInjector::RemoteVirtualAlloc();

	char buf[MAX_PATH]={0};
	GetCurrentDirectory(MAX_PATH,buf);
	strcat_s(buf,MAX_PATH,"\\");
	strcat_s(buf,MAX_PATH,DllPath);
	CDllInjector::SetInjectDll(buf);

	CDllInjector::SyncInject();
	CDllInjector::RemoteVirtualFree();
}
